# -*- coding: utf-8 -*-
'''
Wcrypto Module 
Function used for crypto
implemented with PyCrypto
Make you don't care about Key Size,Block Size
hash your key into the key size and format your data into block size
set Crypto method like this 'AES-CBC-32'
AES is a Cipher
CBC is a Cipher Mode
32 is Key Size
default sets 'AES-CBC-32'
Here is a list of Ciphers and Modes you can used
-----------Ciphers-----------------------
Cipher             Key Size/Block Size
AES                16, 24, or 32 bytes/16 bytes
ARC2               Variable/8 bytes
ARC4               Variable/1 bytes
Blowfish           Variable/8 bytes
CAST               5-16 bytes/8 bytes
DES                8 bytes/8 bytes
DES3 (Triple DES)  16 bytes/8 bytes
IDEA               16 bytes/8 bytes###removed by GAE
RC5                at least 2 bytes/8 bytes###removed by GAE
XOR                1-32bytes/1 bytes
-----------------------------------------
New PyCrypto version 2.1.0alpha1 has removed the IDEA and RC5 modules due to software patents.(I also support)
-----------Modes-------------------------
ECB fastest        0.391s    may be dangerous(unsafe)
CBC middle         0.453s    default and suggest
CFB slowest        5.375s
CTR                unknown   #undocument by PyCrypto
OFB                0.437s    #undocument by PyCrypto
PGP                0.438s    #removed by GAE
-------------------------------------------

1.the cipher which block size is 1 DO NOT have modes
2.RC5 with keysize=1,17 will cause a MemmoryError
3.if you use those Ciphers and Modes which are removed by GAE on GAE will    
     raise a ImportError
Example:
akey=Wcrypto.Wcrypto('AES-CBC-32') #for init 
encrypted_data=akey.encrypt(data,key) #encrypt
dosomethingwith(encrypted_data) #send,store or something else
data=akey.decrypt(encrypted_data,key) #decrypt'''

__author__ = 'redsxz(bronze1man@gmail.com)'
__version__ = '0.4.0'

import hashlib, os

class XOR:
    '''In case No Crypto, wo can also XOR with pure Python'''
    def __init__(self, key):
        self.key = key

    def encrypt(self, data):
        import itertools
        xorsize = 1024
        key = [ord(i) for i in self.key]
        start = 0; end = len(data)
        ss = [None]*(end/xorsize+1); i=0
        ck = itertools.cycle(key)
        while start < end:
            dend = start + xorsize
            dd = [ord(d)^k for d,k in itertools.izip(data[start:dend], ck)]
            ss[i] = ''.join(map(chr, dd)); i+=1
            start = dend
        return ''.join(ss)
    decrypt = encrypt


class Wcrypto:

    def __init__(self, smode='AES-CBC-32'):
        smode = smode.split('-')
        smode = smode + [''] * (4 - len(smode))
        self.cipher = smode[0] if smode[0] else 'AES'
        if self.cipher not in self._ListofCiphers:
            raise ValueError('Invalid cipher: '+self.cipher)
        self.mode = smode[1] if (smode[1] in self._ListofModes) else 'CBC'
        self.blocksize = self._ListofBlockSize[self.cipher]
        if self.blocksize==1: self.mode=''
        try:
            self.keysize = int(smode[2])
        except ValueError:
            self.keysize = 32
        self.hashkey = False if (self.keysize==0 or smode[3]=='0') else True
        if self.keysize != 0:
            if self.cipher not in self._ListofKeySize:
                if self.keysize<=0: self.keysize=32
            elif self.keysize not in self._ListofKeySize[self.cipher]:
                self.keysize = self._ListofKeySize[self.cipher][-1]
        #follow will cause a Memmory Error
        if self.cipher=='RC5' and self.keysize in (1, 17): self.keysize=32
        #try to import Crypto.Cipher.xxxx
        try:
            cipherlib = __import__('Crypto.Cipher.'+self.cipher, fromlist=[''])
            self._newobj = cipherlib.new
            if self.blocksize != 1:
                self._ciphermode = getattr(cipherlib, 'MODE_'+self.mode)
        except ImportError, e:
            if self.cipher=='XOR': self._newobj=XOR
            else: raise e

    def _getcrypto(self, key):
        if not self.hashkey and key=='': return None
        khash = hashlib.sha512(key).digest()
        if self.hashkey: key = khash
        if self.keysize != 0:
            key = (key+khash)[:self.keysize]
        if self.blocksize == 1: return self._newobj(key)
        return self._newobj(key, self._ciphermode, khash[-self.blocksize:])

    def encrypt(self, data, key):
        crypto = self._getcrypto(key)
        if crypto is None: return data
        if self.blocksize != 1:
            patlen = (self.blocksize - len(data) - 1) % self.blocksize
            data = '%s%s%s' % (data, os.urandom(patlen), chr(patlen))
        return crypto.encrypt(data)

    def decrypt(self, data, key):
        crypto = self._getcrypto(key)
        if crypto is None: return data
        data = crypto.decrypt(data)
        if self.blocksize != 1:
            patlen = ord(data[-1])
            data = data[:-(patlen+1)]
        return data

    def realmodestr(self):
        return '%s-%s-%d-%d' % (self.cipher, self.mode, self.keysize, self.hashkey)

    _ListofCiphers = ['AES','ARC2','ARC4','Blowfish','CAST','DES','DES3','IDEA','RC5','XOR']
    _ListofModes = ['ECB','CBC','CFB','OFB','PGP']#Don't know how use CTR
    _ListofKeySize = {'AES':[16,24,32],
                      'DES':[8],
                      'DES3':[16],
                      'IDEA':[16],
                      'CAST':range(5,17)}
    _ListofBlockSize={'AES':16,
                      'ARC2':8,
                      'ARC4':1,
                      'Blowfish':8,
                      'CAST':8,
                      'DES':8,
                      'DES3':8,
                      'IDEA':8,
                      'RC5':8,
                      'XOR':1}

    #_ListofSuggestCipers=['AES','ARC2','ARC4','Blowfish','CAST','DES','DES3']
    #_ListofSuggestModes=['ECB','CBC','CFB','OFB']